
import os
import torch
from torch import nn
import numpy as np
import matplotlib.pyplot as plt

from torch.nn import functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torchvision.utils import make_grid

from tqdm.notebook import tqdm

#Libraries to save images to animations
from PIL import Image as Img, ImageSequence
from IPython.display import Image, display


import helper #Custom library to help with data loading and saving samples

transform = transforms.ToTensor()


device = helper.get_default_device()
print(device)

mps


#trianing_defaults
batch_size = 128
latent_size = 128

real_targets = torch.ones(batch_size,1)
fake_targets = torch.zeros(batch_size, 1,device=device)

#saved directories
sample_dir = 'generated_5'
model_dir = 'saved_models_5'
animation_name = 'animation.gif'

#Fixed Latent
fixed_latent = torch.randn(64, latent_size, 1, 1, device=device)


#Training defaults
lr_d = 0.001
lr_g = 0.0005
epochs = 25


#Label smoothing
real_targets = real_targets - 0.2 + (torch.randn(real_targets.shape) * 0.05)
real_targets = real_targets.to(device)
# print(torch.mean(real_targets),torch.max(real_targets), torch.min(real_targets))


if 'cifar-10-batches-py' in os.listdir('data/'):
    train_ =True
else:
    train_ = False


dataset = datasets.CIFAR10(root = './data',train=train_,download=False, transform = transform)
dataloader = DataLoader(dataset, batch_size=batch_size, drop_last=True)

#Moving data to GPU/MPS
train_dl = helper.DeviceDataLoader(dataloader,device)


generator = nn.Sequential(
 
    #in: latent_size x 1 x 1
    nn.ConvTranspose2d(latent_size, 512, kernel_size=4, stride=1, padding=0, bias=False),
    nn.BatchNorm2d(512),
    nn.ReLU(True),
    #out: 512 x 4 x 4
    
    
    nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(256),
    nn.ReLU(True),
    # out: 256 x 8 x 8

    nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(128),
    nn.ReLU(True),
    # out: 128 x 16 x 16

    nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(64),
    nn.ReLU(True),
    # out: 64 x 32 x 32
    
    nn.ConvTranspose2d(64, 32, kernel_size=1, stride=1, padding=0, bias=False),
    nn.BatchNorm2d(32),
    nn.ReLU(True),
    # out: 32 x 32 x 32

    nn.ConvTranspose2d(32, 3, kernel_size=1, stride=1, padding=0, bias=False),
    nn.BatchNorm2d(3),
    nn.Tanh()
#     out: 3 x 32 x 32
)


discriminator = nn.Sequential(
    # in: 3 x 32 x 32

    nn.Conv2d(3,64,kernel_size=4,stride=2,padding=1,bias=False),
    nn.BatchNorm2d(64),
    nn.LeakyReLU(0.2,inplace=True),

    #second convolutional layer : input = 64*16*16, output = 128*8*8,
    nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(128),
    nn.LeakyReLU(0.2, inplace=True),

    #third convolutional layer : input = 128*8*8, output = 256*4*4,
    nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(256),
    nn.LeakyReLU(0.2, inplace=True),

    #fourth convolutional layer : input = 256*4*4, output = 512*8*8,
    nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
    nn.BatchNorm2d(512),
    nn.LeakyReLU(0.2, inplace=True),
    nn.Conv2d(512, 1, kernel_size=2, stride=1, padding=0, bias=False),
    nn.Flatten(),
    nn.Sigmoid())


# Moving both the models to GPU
discriminator = discriminator.to(device)
generator = generator.to(device)


def train_discriminator(real_images, opt_d):
    ''' Function to train discriminator'''
    # Clear discriminator gradients
    opt_d.zero_grad()

    # Pass real images through discriminator
    real_preds = discriminator(real_images)
    real_loss = F.binary_cross_entropy(real_preds, real_targets)
#     real_score = torch.mean(real_preds).item()
    
    # Generate fake images using the latest generator
    latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
    fake_images = generator(latent)

    # Pass fake images through discriminator
    fake_targets = torch.zeros(fake_images.size(0), 1, device=device)
    fake_preds = discriminator(fake_images)
    fake_loss = F.binary_cross_entropy(fake_preds, fake_targets)
#     fake_score = torch.mean(fake_preds).item()

    # Update discriminator weights
    loss = real_loss + fake_loss
    loss.backward()
    opt_d.step()
    
    return loss.item()


def train_generator(opt_g):
    ''' Function to train discriminator'''
    
    # Clear generator gradients
    opt_g.zero_grad()
    
    # Generate fake images
    latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
    fake_images = generator(latent)
    
    # Pass the samples through the discriminator
    preds = discriminator(fake_images)
    targets = torch.ones(batch_size, 1, device=device)
    loss = F.binary_cross_entropy(preds, targets)
    
    # Update generator weights
    loss.backward()
    opt_g.step()
    
    return loss.item()


# helper.save_samples(0,fixed_latent,generator,sample_dir)


def get_norm(model):
    '''The get_norm function keeps track of exploding/vanishing gradients 
    to help us determine when and if the model is being unstable''' 

    total_norm=0
    for p in model.parameters():
        param_norm = p.grad.data.norm(2)
        total_norm += param_norm.item() ** 2

    total_norm = total_norm ** (1. / 2)
    return total_norm



def fit(epochs, lr_g,lr_d, start_idx=1):
    ''' Main function that starts the training loop for discriminator and generator '''

    torch.cuda.empty_cache()
    
    # Losses & scores
    losses_g = []
    losses_d = []
    
    # Create optimizers
    opt_d = torch.optim.Adam(discriminator.parameters(), lr=lr_d, betas=(0.5, 0.999))
    opt_g = torch.optim.Adam(generator.parameters(), lr=lr_g, betas=(0.5, 0.999))
    
    for epoch in range(epochs):
        for real_images, _ in tqdm(train_dl):
            # Train discriminator
            loss_d = train_discriminator(real_images, opt_d)
            # Train generator
            loss_g = train_generator(opt_g)
            
        # Record losses & scores
        losses_g.append(loss_g)
        losses_d.append(loss_d)
#         real_scores.append(real_score)
#         fake_scores.append(fake_score)
        
        
        
        # Log losses & scores (last batch)
        print("Epoch [{}/{}], Generator Loss: {:.4f}, Discriminator Loss: {:.4f}".format(
            epoch+1, epochs, loss_g, loss_d))
        print(f'Generator Norms : {get_norm(generator)}, Discriminator Norms : {get_norm(discriminator)} ')
        
        
        # Save generated images
        generator_scripted = torch.jit.script(generator.to('cpu'))
        generator_scripted.save(f'{model_dir}/generator_{start_idx+epoch}.pt')
        generator.to('mps')
        
        helper.save_samples(epoch+start_idx, fixed_latent,generator,sample_dir)
    
    return losses_g, losses_d

def save_history(history):
    ''' Function to save losses'''
    
    with open(f'{model_dir}/history_data.csv','a') as csvfile:
        np.savetxt(csvfile, history,delimiter=',')
        
    csvfile.close()


history = fit(10, lr_d,lr_g)
save_history(history)

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [1/10], Generator Loss: 4.6450, Discriminator Loss: 1.0447
Generator Norms : 11.504434838370331, Discriminator Norms : 22.63243333522244 
Saving generated-images-0001.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [2/10], Generator Loss: 3.6011, Discriminator Loss: 1.6514
Generator Norms : 3.745516341642759, Discriminator Norms : 14.151016270239364 
Saving generated-images-0002.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [3/10], Generator Loss: 2.8312, Discriminator Loss: 0.9423
Generator Norms : 7.321329276970922, Discriminator Norms : 27.609685746702002 
Saving generated-images-0003.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [4/10], Generator Loss: 1.8237, Discriminator Loss: 0.9705
Generator Norms : 4.401951817216774, Discriminator Norms : 23.410704490601518 
Saving generated-images-0004.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [5/10], Generator Loss: 2.3200, Discriminator Loss: 1.4835
Generator Norms : 15.927289615517495, Discriminator Norms : 49.78101822169262 
Saving generated-images-0005.png


history2 = fit(epochs, lr_d/4,lr_g/4,start_idx=26)
save_history(history2)

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [1/25], Generator Loss: 4.3875, Discriminator Loss: 0.6116
Generator Norms : 85.02541578469378, Discriminator Norms : 131.85608290255743 
Saving generated-images-0026.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [2/25], Generator Loss: 4.5861, Discriminator Loss: 0.7771
Generator Norms : 140.3354616526763, Discriminator Norms : 184.38387590012297 
Saving generated-images-0027.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [3/25], Generator Loss: 3.1698, Discriminator Loss: 0.6810
Generator Norms : 39.93123663917187, Discriminator Norms : 79.7555986690805 
Saving generated-images-0028.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [4/25], Generator Loss: 0.3635, Discriminator Loss: 2.5925
Generator Norms : 18.659732030112643, Discriminator Norms : 322.67323909635024 
Saving generated-images-0029.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [5/25], Generator Loss: 2.9186, Discriminator Loss: 0.6495
Generator Norms : 67.70791063665826, Discriminator Norms : 103.56779243486459 
Saving generated-images-0030.png


history3 = fit(epochs, lr_d/8,lr_g/8,start_idx=51)
save_history(history3)

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [1/25], Generator Loss: 3.4096, Discriminator Loss: 0.6377
Generator Norms : 68.70809976768112, Discriminator Norms : 115.24125701118851 
Saving generated-images-0051.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [2/25], Generator Loss: 3.5632, Discriminator Loss: 0.5716
Generator Norms : 137.2090685275656, Discriminator Norms : 197.58895639246143 
Saving generated-images-0052.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [3/25], Generator Loss: 3.2193, Discriminator Loss: 0.5829
Generator Norms : 157.89228998482227, Discriminator Norms : 227.23630043614895 
Saving generated-images-0053.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [4/25], Generator Loss: 4.1006, Discriminator Loss: 0.5679
Generator Norms : 124.25948848187944, Discriminator Norms : 222.6435656472736 
Saving generated-images-0054.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [5/25], Generator Loss: 3.6206, Discriminator Loss: 0.6045
Generator Norms : 188.81998807547055, Discriminator Norms : 286.6677652756032 
Saving generated-images-0055.png


history4 = fit(epochs, lr_d/15,lr_g/15,start_idx=76)
save_history(history4)

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [1/25], Generator Loss: 2.8231, Discriminator Loss: 0.5474
Generator Norms : 71.09688462901111, Discriminator Norms : 139.264069362667 
Saving generated-images-0076.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [2/25], Generator Loss: 3.3538, Discriminator Loss: 0.6215
Generator Norms : 114.28085071686013, Discriminator Norms : 272.656780694909 
Saving generated-images-0077.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [3/25], Generator Loss: 3.8461, Discriminator Loss: 0.5240
Generator Norms : 169.48324151260923, Discriminator Norms : 325.12942716655647 
Saving generated-images-0078.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [4/25], Generator Loss: 3.1280, Discriminator Loss: 0.5228
Generator Norms : 94.01124874028197, Discriminator Norms : 177.74481769022603 
Saving generated-images-0079.png

  0%|          | 0/390 [00:00<?, ?it/s]

Epoch [5/25], Generator Loss: 4.0838, Discriminator Loss: 0.5421
Generator Norms : 68.39891487390418, Discriminator Norms : 155.3660512404359 
Saving generated-images-0080.png


with open(f'saved_models_4/history_data.csv') as f:
    history = f.readlines()
    
clean = []
for element in history:
    element = element.strip('\n').split(',')
    temp = [float(x) for x in element]
    clean.append(temp)

gen_index = list(range(0,len(clean),2))
disc_index = list(range(1,len(clean),2))
gen_history = [clean[x] for x in gen_index]
gen_history = [item for sublist in gen_history for item in sublist]
disc_history = [clean[x] for x in disc_index]
disc_history = [item for sublist in disc_history for item in sublist]


plt.figure(figsize=(9,7))
plt.plot(gen_history, label = 'Generator Loss')
plt.plot(disc_history, label = 'Discriminator Loss')
plt.xlabel('Epoch',fontsize=14)
plt.ylabel('Mean loss per epoch',fontsize = 14)
plt.legend()
plt.show()


# Set the file names for your images
file_names = os.listdir(f'{sample_dir}/')
file_names = sorted([f'{sample_dir}/'+x for x in file_names])

# Open each image and append it to the list of frames
frames = []
for file_name in file_names:
    with Img.open(file_name) as image:
        frames.append(image.copy())

# Save the frames as an animated GIF
frames[0].save(f'{animation_name}', format='GIF', append_images=frames[1:], save_all=True, duration=200, loop=0)


# Display the animated GIF
with open('animation.gif','rb') as f:
    display(Image(data=f.read(), format='png',width=420),clear=True)


new_latent = torch.randn(64, latent_size, 1, 1, device=device)

#Loading a random generator from any of the epochs
x = [x for x in os.listdir(f'{model_dir}') if 'generator_' in x and 'pt' in x]
valid_epochs = [int(k.split('_')[1].strip('.pt')) for k in x if int(k.split('_')[1].strip('.pt')) > 30]
generator_rand = np.random.choice(valid_epochs)
final_generator = torch.jit.load(f'{model_dir}/generator_{generator_rand}.pt')

final_generator.to(device)
generated_images = final_generator(fixed_latent)

helper.save_samples(1011,new_latent,final_generator,f'{sample_dir}/')
# Display the generated_image
with open(f'{sample_dir}/generated-images-1011.png','rb') as f:
    display(Image(data=f.read(), format='png',width=400),clear=True,metadata = generator_rand)
    print(f'Generator used : {model_dir}/generator_{generator_rand}.pt')

Generator used : saved_models_5/generator_87.pt

Step 1 : Configuring default values for network training¶

Step 2 : Dataset Handling¶

Step 3 : Defining Generator/Discriminator Architecture¶

Training approach:¶

Step 4 : Plotting the generator and discriminator loss¶

Step 5 : Animation to show different images generated during different epochs¶

Animation plot for generated images in 80 epochs¶

Step 6 : Displaying a new image from any random latent vector¶